#!/usr/bin/env python
# -*- coding:utf-8 -*-

import os, cPickle
import numpy as np
import argparse


def voc_ap(rec, prec, use_voc_metric, use_07_metric):
    """ ap = voc_ap(rec, prec, [use_07_metric])
    Compute VOC AP given precision and recall.
    If use_07_metric is true, uses the
    VOC 07 11 point method (default:False).
    """
    if use_voc_metric:
        if use_07_metric:
            # 11 point metric
            ap = 0.
            for t in np.arange(0., 1.1, 0.1):
                if np.sum(rec >= t) == 0:
                    p = 0
                else:
                    p = np.max(prec[rec >= t])
                ap = ap + p / 11.
        else:
            # correct AP calculation
            # first append sentinel values at the end
            mrec = np.concatenate(([0.], rec, [1.]))
            mpre = np.concatenate(([0.], prec, [0.]))

            # compute the precision envelope
            for i in range(mpre.size - 1, 0, -1):
                mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

            # to calculate area under PR curve, look for points
            # where X axis (recall) changes value
            i = np.where(mrec[1:] != mrec[:-1])[0]

            # and sum (\Delta recall) * prec
            ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    else:
        ap = 2/(1/rec[-1] + 1/prec[-1])

    return ap 


def get_prec_rec(mypath, use_voc_metric=True, use_07_metric=True):
    with open(mypath, 'r') as f:
        pr = cPickle.load(f)
        prec = pr['prec']
        rec = pr['rec']
    prec = np.array(prec)
    rec = np.array(rec)
    ap = voc_ap(rec, prec, use_voc_metric, use_07_metric)

    if not use_voc_metric:
        print('Precision for {} = {:.4f}'.format(os.path.split(mypath)[1][:-7], prec[-1]))
        print('Recall for {} = {:.4f}'.format(os.path.split(mypath)[1][:-7], rec[-1]))

    print('AP for {} = {:.4f}'.format(os.path.split(mypath)[1][:-7], ap))


def parse_args():
    """Parse input arguments."""
    parser = argparse.ArgumentParser(description='Faster R-CNN voc_eval')
    parser.add_argument('--path', dest='pr_path',
                        help='set the pred_file path',
                        type=str)
    parser.add_argument('--not07', dest='metric2007',
                        help='set whether use 2007 metric or not',
                        action='store_false')
    parser.add_argument('--notvoc', dest='voc_metric',
                        help='set whether use voc metric or not',
                        action='store_false')

    myargs = parser.parse_args()
    return myargs


if __name__ == '__main__':
    # --path voc_2007_test/pvanet_frcnn_car_20171212_iter_100K
    # --path voc_2007_train/pvanet_frcnn_car_20171212_iter_100K
    args = parse_args()
    pkldirpath = os.path.join('/home/hrs/installations/pva-faster-rcnn/output/submit_1019', args.pr_path)

    if args.voc_metric:
        print('--- Use VOC 2007 metric? ' + ('Yes' if args.metric2007 else 'No'))
    else:
        print('--- Use VOC metric? ' + ('Yes' if args.voc_metric else 'No'))
    for pklfile in sorted(os.listdir(pkldirpath)):
        if pklfile == 'detections.pkl':
            continue

        get_prec_rec(os.path.join(pkldirpath, pklfile), args.voc_metric, args.metric2007)
